/* ad.c -- atmel downloader
**
** Copyright 2006, Brian Swetland.  All rights reserved.
** See provided LICENSE file or http://frotz.net/LICENSE for details.
*/

#include <stdio.h>
#include <stdlib.h>
#include <unistd.h>
#include <fcntl.h>
#include <string.h>
#include <signal.h>

#include "remote.h"
#include "protocol.h"

static volatile int KILLED = 0;

static void sigint_handler(int n)
{
    signal(SIGINT, sigint_handler);
    KILLED = 1;
}

#define MC_FMR       0xffffff60
#define MC_FCR       0xffffff64
#define MC_FSR       0xffffff68

#define MC_FRDY      0x00000001
#define MC_LOCKE     0x00000004
#define MC_PROGE     0x00000008
#define MC_SECURITY  0x00000010
#define MC_GPNVM0    0x00000100
#define MC_GPNVM1    0x00000200

#define MC_NEBP      0x00000080
#define MC_FWS_1     0x00000000
#define MC_FWS_2     0x00000100
#define MC_FWS_3     0x00000200
#define MC_FWS_4     0x00000300
#define MC_FMCN(n)   (((n) & 0xff) << 16)

#define MC_PAGEN(n)  (((n) & 0x3ff) << 8)
#define MC_KEY       0x5a000000
#define MC_CMD_NONE  0x00000000
#define MC_CMD_WP    0x00000001
#define MC_CMD_SLB   0x00000002
#define MC_CMD_WPL   0x00000003
#define MC_CMD_CLB   0x00000004
#define MC_CMD_EA    0x00000008
#define MC_CMD_SGPB  0x0000000b
#define MC_CMD_CGPB  0x0000000d
#define MC_CMD_SSB   0x0000000f

void *load_file(const char *name, int *sz)
{
    int fd;
    char *data = 0;
    int size;

    fd = open(name, O_RDONLY);
    if(fd < 0) return 0;
    size = lseek(fd, 0, SEEK_END);
    if(size < 0) goto oops;
    if(lseek(fd, 0, SEEK_SET) <0) goto oops;
    if((data = (char*) malloc(size + 1024)) == 0) goto oops;
    if(read(fd, data, size) != size) goto oops;

    *sz = size;
    return (void*) data;

oops:
    close(fd);
    if(data) free(data);
    return 0;
}

static void do_read(char **args)
{
    unsigned addr = strtoul(args[0], 0, 16);
    if(addr & 3) die("unaligned address");
    printf("%08x: %08x\n", addr, read32(addr));
}

static void do_write(char **args)
{
    unsigned addr = strtoul(args[0], 0, 16);
    unsigned val  = strtoul(args[1], 0, 16);
    if(addr & 3) die("unaligned address");
    write32(addr,val);
    printf("%08x: %08x\n", addr, read32(addr));
}

static void download_and_run(char *fn, unsigned addr, unsigned start)
{
    int size;
    int xfer;
    void *data;
    unsigned char *ptr;
    
    if(stop) stop();
    
    data = load_file(fn, &size);
    size = (size + 3) & (~3);

    if(data == 0) {
        printf("could not load '%s'\n", fn);
        return;
    }

    printf("uploading '%s' to 0x%08x (%d bytes)\n", fn, addr, size);
    ptr = (unsigned char*) data;
    while(size > 0){
        xfer = size > 4096 ? 4096 : size;        
        write_words(addr, ptr, xfer);
        addr += xfer;
        ptr += xfer;
        size -= xfer;
    }
    printf("running code at 0x%08x\n", start);
    jump_to_addr(start);
}

static void do_run(char **args)
{
    download_and_run(args[0], 0x202000, 0x202000);
}

static void do_boot(char **args)
{
    download_and_run(args[0], 0x202000, 0x202014);
}

static void do_stop(char **args)
{
    if(stop) stop();
}

static void do_resume(char **args)
{
    if(resume) resume();
}

static void do_regs(char **args)
{
    if(read_regs) {
        unsigned R[17];
        read_regs(R);

        printf(" R0 %08x  R1 %08x  R2 %08x  R3 %08x\n",
               R[0], R[1], R[2], R[3]);
        printf(" R4 %08x  R5 %08x  R6 %08x  R7 %08x\n",
               R[4], R[5], R[6], R[7]);
        printf(" R8 %08x  R9 %08x R10 %08x R11 %08x  CPSR\n",
               R[8], R[9], R[10], R[11]);
        printf("R12 %08x  SP %08x  LR %08x  PC %08x  %08x\n",
               R[12], R[13], R[14], R[15], R[16]);
    }
}

/* BUG: adjust this based on product ID */
#define FLASH_BOTTOM     0x00100000
#define FLASH_SIZE       (64*1024)
#define FLASH_TOP        (FLASH_BOTTOM + FLASH_SIZE)
#define FLASH_PAGE_SIZE  128
#define FLASH_PAGE_MASK  (FLASH_PAGE_SIZE - 1)
#define FLASH_LOCK_BITS  16

static void flash_command(unsigned cmd, unsigned pagen)
{
    unsigned n;

    while(!(read32(MC_FSR) & MC_FRDY));
    write32(MC_FCR, MC_PAGEN(pagen) | cmd | MC_KEY);
    for(;;) {
        n = read32(MC_FSR);
        if(n & MC_PROGE) die("flash programming error");
        if(n & MC_LOCKE) die("flash lock error");
        if(n & MC_FRDY) break;
    }
}

static void do_unlock(char **args)
{
    unsigned n, lockbits;

    lockbits = read32(MC_FSR) >> 16;
    
        /* BUG: hardcoded timing for 18MHz */
    write32(MC_FMR, 0x00050100);

    for(n = 0; n < FLASH_LOCK_BITS; n++){
        if(lockbits & 1) {
            printf("unlocking lock bit #%d\n", n);
            flash_command(MC_CMD_CLB,
                          (n * (FLASH_SIZE / FLASH_LOCK_BITS)) / FLASH_PAGE_SIZE);
        }
        lockbits >>= 1;
    }
}

static void do_patch4k(char **args)
{
    unsigned data[FLASH_PAGE_SIZE / 4];

    read_words(0, data, FLASH_PAGE_SIZE);
    data[0] = 0xea0007fe; // relative jump from 0 to 0x2000
    write_words(0x00100000, data, FLASH_PAGE_SIZE);
    flash_command(MC_CMD_WP, 0);
}

static void do_flash(char **args)
{
    int size;
    void *data;
    unsigned addr;
    unsigned char *ptr;
    unsigned char readback[512];
    
    addr = strtoul(args[1], 0, 16);
    data = load_file(args[0], &size);
    
    if(data == 0) die("cannot load file");
    size = (size + FLASH_PAGE_MASK) & (~FLASH_PAGE_MASK);
    
    if(addr & 127) die("not page aligned");
    if((addr < FLASH_BOTTOM) || (addr > FLASH_TOP)) {
        die("address out of range");
    }
    if((addr + size) > FLASH_TOP) {
        die("image will not fit");
    }

        /* BUG: hardcoded timing for 18MHz */
    write32(MC_FMR, 0x00340000);
    
    printf("writing %d bytes to flash @ 0x%08x\n", size, addr);
    ptr = (unsigned char*) data;
    while(size > 0){
        write_words(addr, ptr, FLASH_PAGE_SIZE);
        flash_command(MC_CMD_WP, (addr & (FLASH_SIZE - 1)) / FLASH_PAGE_SIZE);
        read_words(addr, readback, FLASH_PAGE_SIZE);
        if(memcmp(readback, ptr, FLASH_PAGE_SIZE)) {
            die("verify error");
        }
        size -= FLASH_PAGE_SIZE;
        addr += FLASH_PAGE_SIZE;
        ptr += FLASH_PAGE_SIZE;
    }
}

static void do_dump(char **args)
{
    unsigned addr = strtoul(args[0], 0, 16);
    unsigned words[256];
    unsigned n;
    if(addr & 3) die("unaligned address");
    read_words(addr, words, 256);
    for(n = 0; n < 256; n += 16){
        printf("%08x: %08x %08x %08x %08x\n",
               addr + n,
               words[n/4], words[n/4+1],
               words[n/4+2], words[n/4+3]);
    }
}

static void do_listen(char **args)
{
    debug_bridge db;
    unsigned n;

    while(!KILLED){
        read_words(DEBUG_BRIDGE_ADDR, &db, 8);
        if(db.handshake != DEBUG_HANDSHAKE){
            usleep(100);
            continue;
        }
        read_words(DEBUG_BRIDGE_ADDR + 8, db.data, DEBUG_MAX_DATA);
        n = db.count > DEBUG_MAX_DATA ? DEBUG_MAX_DATA : db.count;
        db.handshake = 0;
        db.count = 0;
        write_words(DEBUG_BRIDGE_ADDR, &db, 8);

        if((n == 7) && (memcmp(db.data, "@done@\n", 7) == 0)) {
            break;
        }
        fwrite(db.data, n, 1, stdout);
        fflush(stdout);
    }
}

static void do_setclk(char **args)
{
    int timeout;
    
    write32(0xfffffc20, 0x00000000); /* reset MOSCS */
    write32(0xfffffc20, 0x00004001); /* enable main oscillator */

    for(timeout = 0; timeout < 100; timeout++){
        if(read32(0xfffffc68) & 1) goto okay1;
    }
    die("setclk: timeout");

okay1:
    write32(0xfffffc30, 0x00000001); /* MCK = main osc / 1 */
    
    for(timeout = 0; timeout < 100; timeout++){
        if(read32(0xfffffc68) & 8) return; 
    }
    die("setclk: timeout");
}

static void do_jtag(char **args)
{
    unsigned n = read32(0xee000000);
    if(n != 0x4741544a) {
        die("bootloader is non-JTAG version");
    }
}

static void do_version(char **args)
{
    unsigned signature[5];

    read_words(0xee000000, signature, 16);
    signature[4] = 0;
    printf("version: %s\n", (char*) signature);
}

static void do_help(char **args);

struct 
{
    const char *name;
    const char *usage;
    int args;
    void (*func)(char **args);
} cmdlist[] = {
    { "read",   "<addr>            read uint32",                1, do_read },
    { "rd",     "<addr>            read uint32",                1, do_read },
    { "r",      "<addr>            read uint32",                1, do_read },
    { "write",  "<addr> <value>    write uint32",               2, do_write },
    { "wr",     "<addr> <value>    write uint32",               2, do_write },
    { "w",      "<addr> <value>    write uint32",               2, do_write },
    { "run",    "<file>            upload to 0x202000 and run", 1, do_run },
    { "boot",   "<file>            upload and run bootloader",  1, do_boot },
    { "unlock", "                  unlock flash",               0, do_unlock },
    { "flash",  "<file> <addr>     write file to flash",        2, do_flash },
    { "setclk", "                  prepare for jtag flash",     0, do_setclk },
    { "jtag",   "                  fail if not a jtag target",  0, do_jtag },
    { "dump",   "<addr>            dump 64 uint32s",            1, do_dump },
//    { "patch4k","hack sam-ba boot for a second stage at 0x2000",0, do_patch4k },
    { "stop",   "                  suspend non-debugger code",  0, do_stop },
    { "resume", "                  resume non-debugger code",   0, do_resume },
    { "regs",   "                  dump registers",             0, do_regs },
    { "listen", "                  poll for dprintf chatter",   0, do_listen },
    { "version","                  print bootloader version",   0, do_version},
    { "help",   0,                                              0, do_help },
    { 0, 0, 0, 0 },
};
      
static void do_help(char **args)
{
    int i;
    
    for(i = 0; cmdlist[i].usage; i++) {
        printf("%-12s %s\n", cmdlist[i].name, cmdlist[i].usage);
    }
}

static void open_remote(int wait)
{
    for(;;) {
        if(open_remote_debugger() == 0) return;
        if(open_remote_samba() == 0) {
            fprintf(stderr,"WARNING: using slow/unreliable SAM-BA protocol\n");
            return;
        }
        switch(wait) {
        case 0:
            die("cannot find AT91SAM7 device");
        case 1:
            wait = 2;
            fprintf(stderr,"Waiting for AT91SAM7 to connect\n");
        default:
            usleep(100);
	}
    }
}

int main(int argc, char **argv)
{
    int i;

    signal(SIGINT, sigint_handler);
    
    if((argc > 1) && (!strcmp(argv[1], "wait"))){
        open_remote(1);
	argc--;
	argv++;
    } else {
	open_remote(0);
    }

    argc--;
    argv++;

    if(argc == 0) {
        do_help(0);
        return 0;
    }
    
    while(argc-- > 0){
        char *cmd = *argv++;
        
        for(i = 0; cmdlist[i].name; i++){
            if(!strcmp(cmd, cmdlist[i].name)){
                if(argc < cmdlist[i].args) {
                    die("not enough arguments");
                }                    

                cmdlist[i].func(argv);
                argv += cmdlist[i].args;
                argc -= cmdlist[i].args;
                break;
            }
        }
        if(cmdlist[i].name == 0) {
            die("unknown command");
        }
    }

    close_remote();
    return 0;
}

void die(const char *reason)
{
    fprintf(stderr,"-- error: %s --\n", reason);
    exit(1);
}


void (*close_remote)() = 0;
void (*read_bytes)(unsigned addr, void *ptr, unsigned count) = 0;
void (*read_words)(unsigned addr, void *ptr, unsigned count) = 0;
void (*write_bytes)(unsigned addr, void *ptr, unsigned count) = 0;
void (*write_words)(unsigned addr, void *ptr, unsigned count) = 0;
void (*jump_to_addr)(unsigned addr) = 0;
unsigned (*read32)(unsigned addr) = 0;
unsigned (*read8)(unsigned addr) = 0;
void (*write32)(unsigned addr, unsigned value) = 0;
void (*write8)(unsigned addr, unsigned value) = 0;
void (*stop)(void) = 0;
void (*resume)(void) = 0;
void (*read_regs)(unsigned *r) = 0;

